#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
@author coldwind
"""

import tensorflow as tf


class TextRNN:
    def __init__(self,
                 model_type,
                 sequence_length,
                 number_class,
                 word_embedding,
                 vocab_size,
                 embedding_size,
                 number_layers,
                 number_units,
                 l2_reg_lambda=0.0,
                 model='lstm'):

        self.input_x = tf.placeholder(shape=[None, sequence_length], dtype=tf.int32, name='input_x')
        self.input_y = tf.placeholder(shape=[None, number_class], dtype=tf.float32, name='input_y')
        self.dropout_keep_prob = tf.placeholder(dtype=tf.float32, name='dropout_keep_prob')
        self.learning_rate = tf.placeholder(dtype=tf.float32, name='learning_rate')

        self.l2_loss = tf.constant(0.0, name='l2_loss')

        # embedding
        with tf.device('/cpu:0'), tf.name_scope('embedding'):
            if word_embedding:
                self.Embedding = tf.Variable(initial_value=word_embedding, dtype=tf.float32, name='word_embedding')
            else:
                self.Embedding = tf.get_variable(
                    initializer=tf.random_uniform(shape=[vocab_size, embedding_size], minval=-1, maxval=1),
                    dtype=tf.float32, name='word_embedding')

            self.embedded_words = tf.nn.embedding_lookup(params=self.Embedding, ids=self.input_x)

        # Create a recurrent layer for each rnn layer
        with tf.name_scope(model):
            if 'lstm' == model:
                cell_func = tf.nn.rnn_cell.LSTMCell
            elif 'gru' == model:
                cell_func = tf.nn.rnn_cell.GRUCell
            elif 'rnn' == model:
                cell_func = tf.nn.rnn_cell.BasicRNNCell

            def get_cell(num_units):
                temp_cell = cell_func(num_units, state_is_tuple=True)
                temp_cell = tf.nn.rnn_cell.DropoutWrapper(cell=temp_cell, output_keep_prob=self.dropout_keep_prob)
                return temp_cell

            cell = tf.nn.rnn_cell.MultiRNNCell(cells=[get_cell(number_units) for _ in range(number_layers)])

            # tf.nn.dynamic_rnn(
            #     cell,
            #     inputs,
            #     sequence_length=None,
            #     initial_state=None,
            #     dtype=None,
            #     parallel_iterations=None,
            #     swap_memory=False,
            #     time_major=False,
            #     scope=None
            # )
            # outputs: [batch_size, sequence_length, number_units]
            # last_state: [batch_size, number_units]
            outputs, last_state = tf.nn.dynamic_rnn(cell=cell, inputs=self.embedded_words, dtype=tf.float32)

            # 第一种: 取number_units 每一维度的平均值(可以进行其他操作)
            # [batch_size, number_units]
            self.output = tf.reduce_mean(outputs, axis=1, name='output')

            # 第二种: 将所有时刻的输出进行拼接
            # [batch_size, number_units * sequence_length]
            # self.output = tf.reshape(outputs, shape=[-1, number_units * sequence_length], name='output')

        # dropout
        with tf.name_scope('dropout'):
            self.h_dropout = tf.nn.dropout(self.output, keep_prob=self.dropout_keep_prob, name='dropout')

        # final-output
        with tf.name_scope('final-output'):
            in_dim = self.h_dropout.shape[-1].value
            w = tf.get_variable(shape=[in_dim, number_class], initializer=tf.contrib.layers.xavier_initializer(),
                                dtype=tf.float32, name='W')
            b = tf.Variable(tf.constant(0.1, shape=[number_class]), dtype=tf.float32, name='b')

            self.l2_loss += tf.nn.l2_loss(w)
            self.l2_loss += tf.nn.l2_loss(b)

            self.scores = tf.nn.xw_plus_b(self.h_dropout, weights=w, biases=b, name='scores')
            if 'cls' == model_type:
                self.predictions = tf.argmax(self.scores, axis=1, name='predictions')  # dtype 为 int64
            elif 'reg' == model_type:
                self.predictions = tf.reduce_max(self.scores, axis=1, name='predictions')
                self.predictions = tf.expand_dims(self.predictions, -1)

        # loss
        with tf.name_scope('loss'):
            if 'cls' == model_type:
                losses = tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.input_y, logits=self.scores)
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda
            elif 'reg' == model_type:
                losses = tf.sqrt(tf.losses.mean_squared_error(labels=self.input_y, predictions=self.predictions))
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda

        with tf.name_scope('accuracy'):
            if 'cls' == model_type:
                # correct_predictions:[batch_size, ] dtype 为boolean, 不能直接reduce_mean
                correct_predictions = tf.equal(tf.argmax(self.input_y, axis=1), self.predictions)
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, dtype=tf.float32), name='accuracy')
            elif 'reg' == model_type:
                self.accuracy = tf.constant(0.0, name='accuracy')


if __name__ == '__main__':
    text_rnn = TextRNN(model_type='cls',
                       sequence_length=10,
                       number_class=2,
                       word_embedding=None,
                       vocab_size=100,
                       embedding_size=256,
                       number_layers=2,
                       number_units=128)
